#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 1/12/18 2:50 PM
@desc: predict image using pd
"""
import cv2
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from crowdcounting.data.tools import get_density_map
from utils.basic import get_file_name
from utils.io import read_text_file


def predict_image(image, points, checkpoint):
    size = 256
    shape = (size, size)
    img = cv2.imread(image)
    img = cv2.resize(img, shape)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32)
    img[:, :, 0] -= 80.197326
    img[:, :, 1] -= 75.26667
    img[:, :, 2] -= 72.822747
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        meta_graph_def = tf.saved_model.loader.load(sess, ['frozen_model'], checkpoint)
        signature = meta_graph_def.signature_def
        image_tensor_name = signature['predict_images'].inputs['images'].name
        pre_density_map_name = signature['predict_images'].outputs['density_map'].name
        classify_result_name = signature['predict_images'].outputs['class'].name
        image_tensor = sess.graph.get_tensor_by_name(image_tensor_name)
        pre_density_map = sess.graph.get_tensor_by_name(pre_density_map_name)
        classify_result = sess.graph.get_tensor_by_name(classify_result_name)
        img_np = np.expand_dims(img, axis=0)
        pre_density_map, classify_result = sess.run(
            [pre_density_map, classify_result],
            feed_dict={
                image_tensor: img_np
            })
    points = read_text_file(points)
    points = list(map(lambda x: x.split(' '), points))
    points = list(map(lambda x: [float(x[0]), float(x[1])], points))
    points_array = np.asarray(points) * size
    points_array = points_array.round().astype(np.int32)
    density_map = get_density_map(shape, points_array, kernel_size=15, sigma=4)
    plt.figure(figsize=(22, 10), dpi=120)
    plt.subplot(122)
    plt.title('predict class:{}, number: {}'.format(np.argmax(classify_result), np.sum(pre_density_map)))
    plt.imshow(np.squeeze(pre_density_map), interpolation='none')
    plt.subplot(121)
    bins = 88 / 24
    cls = int((np.sum(density_map) - 24) / bins)
    plt.title('ground class:{}, number: {}'.format(cls, np.sum(density_map)))
    plt.imshow(density_map, interpolation='none')
    plt.savefig('D:\\workspace\\data\\val\\result\\{}.jpg'.format(get_file_name(image)))


if __name__ == '__main__':
    # predict_image('/home/admins/repos/hightensity/val/0000/00000009/edge_24_160_480.jpg',
    #               '/home/admins/repos/hightensity/val/0000/00000009/edge_24_160_480.txt',
    #               '/home/admins/repos/hightensity/model/export/2132/')
    predict_image('D:/workspace/data/val/0000/00000009/edge_18_80_480.jpg',
                  'D:/workspace/data/val/0000/00000009/edge_18_80_480.txt',
                  'D:/workspace/data/val/model/export1/73205/')
